"""
模型模块
包含所有深度学习模型的定义
"""
from .baseline_cnn import BaselineCNNPredictor
from .deep_cnn import DeepCNNPredictor
from .gcnet import GCNetVolumePredictor

__all__ = ['BaselineCNNPredictor', 'DeepCNNPredictor', 'GCNetVolumePredictor']


def build_model(config):
    """
    根据配置构建模型
    
    参数:
        config: 配置对象，包含 model_type, num_images, num_features 等属性
    
    返回:
        model: 构建好的模型实例
    
    示例:
        >>> from configs.default import DefaultConfig
        >>> config = DefaultConfig(model_type='gcnet')
        >>> model = build_model(config)
    """
    model_type = config.model_type.lower()
    
    if model_type == 'baseline_cnn':
        model = BaselineCNNPredictor(
            num_images=config.num_images,
            num_features=config.num_features
        )
        print(f"✓ Built Baseline CNN model")
        
    elif model_type == 'deep_cnn':
        model = DeepCNNPredictor(
            num_images=config.num_images,
            num_features=config.num_features
        )
        print(f"✓ Built Deep CNN model")
        
    elif model_type == 'gcnet':
        model = GCNetVolumePredictor(
            num_images=config.num_images,
            num_features=config.num_features
        )
        print(f"✓ Built GCNet model")
        
    else:
        raise ValueError(f"Unknown model type: {model_type}. "
                        f"Available: 'baseline_cnn', 'deep_cnn', 'gcnet'")
    
    return model


def get_model_info(model_type):
    """
    获取模型信息
    
    参数:
        model_type: 模型类型字符串
    
    返回:
        info: 包含模型信息的字典
    """
    info_dict = {
        'baseline_cnn': {
            'name': 'Baseline CNN',
            'description': '简单的4层CNN + 全连接层',
            'params': '~1.2M',
            'complexity': 'Low',
            'training_time': 'Fast',
            'best_for': '快速实验和基线对比'
        },
        'deep_cnn': {
            'name': 'Deep CNN',
            'description': 'VGG-like深度卷积网络',
            'params': '~5.8M',
            'complexity': 'Medium',
            'training_time': 'Medium',
            'best_for': '平衡性能和速度'
        },
        'gcnet': {
            'name': 'GCNet',
            'description': '全局上下文网络（Global Context Network）',
            'params': '~43.7M',
            'complexity': 'High',
            'training_time': 'Slow',
            'best_for': '追求最佳性能'
        }
    }
    
    return info_dict.get(model_type.lower(), None)


def print_model_comparison():
    """打印所有模型的对比信息"""
    print("\n" + "="*80)
    print("Model Comparison".center(80))
    print("="*80)
    
    for model_type in ['baseline_cnn', 'deep_cnn', 'gcnet']:
        info = get_model_info(model_type)
        print(f"\n【{info['name']}】")
        print(f"  Description: {info['description']}")
        print(f"  Parameters: {info['params']}")
        print(f"  Complexity: {info['complexity']}")
        print(f"  Training Time: {info['training_time']}")
        print(f"  Best For: {info['best_for']}")
    
    print("\n" + "="*80 + "\n")